import joblib
import numpy as np
import pandas as pd
from main import label_encoder

if __name__ == '__main__':
    label_encoder_file = 'label_encoder.pkl'
    model_file = 'random_forest_model.pkl'
    feature_names_file = 'feature_names.pkl'
    random_forest = joblib.load(model_file)
    le_dict = joblib.load(label_encoder_file)
    feature_names = joblib.load(feature_names_file)

    df = pd.read_csv('/tmp/test_data.csv.label.csv')
    df = df.drop(columns=[col for col in
                          ['id', 'THREAT_TIME', 'SIP', 'S_PORT', 'DIP', 'D_PORT', 'XFF_IP', 'PROTOCOL', 'DENY_METHOD',
                           'THREAT_SUMMARY', 'SEVERITY'] if col in df.columns])
    df.replace([np.inf, -np.inf, np.nan], -1, inplace=True)
    df, le_dict = label_encoder(df, le_dict)

    # 对齐特征列
    missing_features = [feature for feature in feature_names if feature not in df.columns]
    for feature in missing_features:
        df[feature] = -1
    df = df[feature_names]
    predictions = random_forest.predict(df)
    print(predictions)
    print(pd.Series(predictions).value_counts())
